import numpy as np
import scipy.io as scio
import yaml
import configparser
from sklearn.model_selection import train_test_split


def square_lr_ud_flip(data_path: str, data_num: int):
    """
    :param data_path: npy文件路径，coding_array
    :param data_num: 加载的数据量
    :return: 向量重整为矩阵，并且上下左右对称
    """
    # 加载文件
    coding_array_all = np.load(data_path)
    coding_array = coding_array_all[0:data_num - 1, :]
    # 将 n*m_1dim 的数组转换为 n*m_2dim*m_2dim 的数组
    reshaped_array = coding_array.reshape((data_num - 1, 5, 5))
    # 左右对称处理
    reshaped_array = np.concatenate([reshaped_array, np.flip(reshaped_array, axis=2)], axis=2)

    # 上下对称处理
    reshaped_array = np.concatenate([reshaped_array, np.flip(reshaped_array, axis=1)], axis=1)

    return reshaped_array


def read_mat_file_to_npy(mat_path: int, data_name: str, data_num: int):
    """
    :param mat_path: mat文件路径
    :param data_name: 变量名称，struct名称
    :param data_num: 加载数据量
    :return: 整合为 n*10 的 nparray 的数据
    """
    mat_data = scio.loadmat(mat_path)[data_name]
    mat_data = np.transpose(mat_data)

    # 定义一个包含1700行和10列的二维列表
    re_data = [[0] * 10 for _ in range(data_num)]

    for i in range(0, 1700):
        for j in range(0, 10):
            re_data[i][j] = data = mat_data[i, 0][j][0]
    all_data = np.array(re_data, dtype=object)
    # np.save('../dataset/all_1700data.npy', alldata) # 数据保存

    return all_data


def load_x_y():
    yaml_file = r'C:\Users\22965\OneDrive - mail.ecust.edu.cn\桌面\Python Code\myVAE\VAE-Bi\config\vae.yaml'
    with open(yaml_file, encoding='utf-8') as f:
        config = yaml.load(stream=f.read(), Loader=yaml.FullLoader)
        data_path = config['data_loader']['path']
        data_size = config['data_loader']['size']
        # 加载数据总集
        coding_array = square_lr_ud_flip(data_path, data_size)
        # 训练集验证集分离
        y = np.zeros((coding_array.shape[0],1))  # 制造一份假的目标数据
        X_train, X_test, y_train, y_test = train_test_split(coding_array, y, test_size=0.2, random_state=0)
    return X_train, X_test


if __name__ == "__main__":
    x_train, x_test = load_x_y() # 需要注意的是，由于运行文件的路径不同，在设置目录时也可能需要改变上下级关系
